module test_mod
  use spark_cubed_sphere
  use params_mod
  use dycore_test
  use static_mod
  use state_mod
  use diag_mod
  use ghost_interp_mod
  use topo_mod
  use vertical_coord_mod
  use met_preproc_mod
  use physics_const_mod
  use hope_recon_mod

  implicit none

  !private
  
  public test_setup
  public test_adv_update_wind
  
  real(8) case_omega
  
contains

  subroutine test_setup(case_id, static, state, diag)

    integer, intent(in) :: case_id
    type(static_type), intent(inout) :: static
    type(state_type ), intent(inout) :: state
    type(diag_type  ), intent(inout) :: diag

    integer idom, ims, ime, jms, jme, kms, kme
    integer       ids, ide, jds, jde, kds, kde
    integer npt, pc, pvs, pve, pes(6), pee(6), pqs, pqe, pgs, pge
    real(8) u, v, w
    real(8), allocatable, dimension(:) :: gd, ght, uc, vc, wc
    real(8), allocatable, dimension(:) :: rho, pt, mr
    real(8)                 :: Jh, JhQ(ncq)
    real(8), dimension(3  ) :: omega
    real(8), dimension(2,2) :: iAh
    real(8) :: time_in_seconds
    integer itime
    integer i, j, k, p
    
    character(500) :: high_field_file, surface_field_file, invariant_field_file
    
    if(model_type==0)then
      associate (mesh => state%mesh)
      call mesh%get_params(npt=npt, pc=pc, pvs=pvs, pve=pve, pes=pes, pee=pee, pqs=pqs, pqe=pqe, pgs=pgs, pge=pge)
      allocate(gd(npt), ght(npt), uc(npt), vc(npt))
      do idom = 1, mesh%ndom
        if (mesh%domain(idom)%initialized) then
          associate (ghost    => mesh%domain(idom)%ghost      , &
                     lon      => mesh%domain(idom)%lon        , &
                     lat      => mesh%domain(idom)%lat        , &
                     r        => mesh%domain(idom)%r          , &
                     q        => state%domain(idom)%q         , &
                     uu       => diag%domain(idom)%uc         , &
                     ucL      => diag%domain(idom)%ucL        , &
                     ucR      => diag%domain(idom)%ucR        , &
                     ucB      => diag%domain(idom)%ucB        , &
                     ucT      => diag%domain(idom)%ucT        , &
                     ucQ      => diag%domain(idom)%ucQ        , &
                     vv       => diag%domain(idom)%vc         , &
                     vcL      => diag%domain(idom)%vcL        , &
                     vcR      => diag%domain(idom)%vcR        , &
                     vcB      => diag%domain(idom)%vcB        , &
                     vcT      => diag%domain(idom)%vcT        , &
                     vcQ      => diag%domain(idom)%vcQ        , &
                     ghs      => static%domain(idom)%ghs      , &
                     ghsC     => static%domain(idom)%ghsC     , &
                     ghsL     => static%domain(idom)%ghsL     , &
                     ghsR     => static%domain(idom)%ghsR     , &
                     ghsB     => static%domain(idom)%ghsB     , &
                     ghsT     => static%domain(idom)%ghsT     , &
                     ghsQ     => static%domain(idom)%ghsQ     , &
                     gravityQ => static%domain(idom)%gravityQ   )
          call mesh%get_params(idom=idom, ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme)
          time_in_seconds = 0
          ! Set static data
          !$OMP PARALLEL DO PRIVATE(p) COLLAPSE(3)
          do k = kms, kme
            do j = jms, jme
              do i = ims, ime
                ! Set gravity, ghs and coriolis on center point
                p = pc
                call adv_case(case_id, time_in_seconds, lon(pc,i,j), lat(pc,i,j), ghs=ghsC(i,j), radi=r(p,i,j,k))
                
                ! Set fields on edge quadrature points.
                do p = pes(1), pee(1)
                  call adv_case(case_id, time_in_seconds, lon(p,i,j), lat(p,i,j), ghs=ghsL(p,i,j), radi=r(p,i,j,k))
                end do
                do p = pes(2), pee(2)
                  call adv_case(case_id, time_in_seconds, lon(p,i,j), lat(p,i,j), ghs=ghsR(p,i,j), radi=r(p,i,j,k))
                end do
                do p = pes(3), pee(3)
                  call adv_case(case_id, time_in_seconds, lon(p,i,j), lat(p,i,j), ghs=ghsB(p,i,j), radi=r(p,i,j,k))
                end do
                do p = pes(4), pee(4)
                  call adv_case(case_id, time_in_seconds, lon(p,i,j), lat(p,i,j), ghs=ghsT(p,i,j), radi=r(p,i,j,k))
                end do
                do p = pes(5), pee(5)
                  r(p,i,j,k) = earth_radius
                end do
                do p = pes(6), pee(6)
                  r(p,i,j,k) = earth_radius
                end do
                
                ! Set fields on cell quadrature points.
                do p = pqs, pqe
                  call adv_case(case_id, time_in_seconds, lon(p,i,j), lat(p,i,j), ghs=ghsQ(p,i,j), g=gravityQ(p,i,j,k), radi=r(p,i,j,k))
                end do
                
                ! Set fields on corner points.
                do p = pvs, pve
                  call adv_case(case_id, time_in_seconds, lon(p,i,j), lat(p,i,j), radi=r(p,i,j,k))
                enddo
                
                ! Set fields on ghost points.
                do p = 1, ghost(i,j)%ngp
                  call adv_case(case_id, time_in_seconds, ghost(i,j)%lon(p), ghost(i,j)%lat(p), radi=ghost(i,j)%r(p,k))
                enddo
              end do
            end do
          enddo
          !$OMP END PARALLEL DO
          
          call mesh%domain(idom)%set_hor_metrics
          
          ! Set fields
          !$OMP PARALLEL DO PRIVATE(p,u,v,ght,uc,vc,gd,Jh,JhQ,iAh) COLLAPSE(3)
          do k = kms, kme
            do j = jms, jme
              do i = ims, ime
                do p = 1, npt
                  iAh = mesh%get_dom_iAh(idom,p,i,j,k)
                  call adv_case(case_id, time_in_seconds, lon(p,i,j), lat(p,i,j), u, v, ght(p))
                  call spherev_to_contrav(iAh, u, v, uc(p), vc(p))
                end do
                
                ! Set ghs on edge quadrature points.
                do p = pes(1), pee(1)
                  ucL(p,i,j,k) = uc(p)
                  vcL(p,i,j,k) = vc(p)
                end do
                do p = pes(2), pee(2)
                  ucR(p,i,j,k) = uc(p)
                  vcR(p,i,j,k) = vc(p)
                end do
                do p = pes(3), pee(3)
                  ucB(p,i,j,k) = uc(p)
                  vcB(p,i,j,k) = vc(p)
                end do
                do p = pes(4), pee(4)
                  ucT(p,i,j,k) = uc(p)
                  vcT(p,i,j,k) = vc(p)
                end do
                
                ! Set initial values on cell quadrature points.
                do p = pqs, pqe
                  ucQ(p,i,j,k) = uc(p)
                  vcQ(p,i,j,k) = vc(p)
                end do
                
                ! Sum quadrature points to get cell integrated average values.
                do p = pqs, pqe
                  JhQ(p-pqs+1) = mesh%get_dom_Jh(idom,p,i,j,k)
                enddo
                ghs(i,j) = mesh%domain(idom)%cell_quad(ghsQ(:,i,j))
                gd(pqs:pqe) = ( ght(pqs:pqe) - ghsQ(:,i,j) ) * JhQ
                uu(i,j,k  ) = mesh%domain(idom)%cell_quad(ucQ(:,i,j,k))
                vv(i,j,k  ) = mesh%domain(idom)%cell_quad(vcQ(:,i,j,k))
                q (i,j,1,1) = mesh%domain(idom)%cell_quad(gd (pqs:pqe))
              end do
            end do
          enddo
          !$OMP END PARALLEL DO
          end associate
        end if
      end do
      deallocate(gd, ght, uc, vc)
      end associate
    elseif(model_type==1)then
      associate (mesh => state%mesh)
      call mesh%get_params(npt=npt, pc=pc, pvs=pvs, pve=pve, pes=pes, pee=pee, pqs=pqs, pqe=pqe, pgs=pgs, pge=pge)
      allocate(gd(pqs:pqe), ght(pqs:pqe), uc(pqs:pqe), vc(pqs:pqe))
      do idom = 1, mesh%ndom
        if (mesh%domain(idom)%initialized) then
          associate (ghost    => mesh%domain(idom)%ghost      , &
                     lon      => mesh%domain(idom)%lon        , &
                     lat      => mesh%domain(idom)%lat        , &
                     r        => mesh%domain(idom)%r          , &
                     q        => state%domain(idom)%q         , &
                     qL       => diag%domain(idom)%qL         , &
                     qR       => diag%domain(idom)%qR         , &
                     qB       => diag%domain(idom)%qB         , &
                     qT       => diag%domain(idom)%qT         , &
                     qQ       => diag%domain(idom)%qQ         , &  
                     ghs      => static%domain(idom)%ghs      , &
                     ghsC     => static%domain(idom)%ghsC     , &
                     ghsL     => static%domain(idom)%ghsL     , &
                     ghsR     => static%domain(idom)%ghsR     , &
                     ghsB     => static%domain(idom)%ghsB     , &
                     ghsT     => static%domain(idom)%ghsT     , &
                     ghsQ     => static%domain(idom)%ghsQ     , &
                     gravityQ => static%domain(idom)%gravityQ , &
                     fC       => static%domain(idom)%coriolisC, &  
                     fQ       => static%domain(idom)%coriolisQ  )
          call mesh%get_params(idom=idom, ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme)
          ! Set static data
          !$OMP PARALLEL DO PRIVATE(p) COLLAPSE(3)
          do k = kms, kme
            do j = jms, jme
              do i = ims, ime
                ! Set gravity, ghs and coriolis on center point
                p = pc
                call sw_case(case_id, lon(pc,i,j), lat(pc,i,j), ghs=ghsC(i,j), f=fC(i,j), radi=r(p,i,j,k))
                
                ! Set gravity, ghs and coriolis on edge quadrature points.
                do p = pes(1), pee(1)
                  call sw_case(case_id, lon(p,i,j), lat(p,i,j), ghs=ghsL(p,i,j), radi=r(p,i,j,k))
                end do
                do p = pes(2), pee(2)
                  call sw_case(case_id, lon(p,i,j), lat(p,i,j), ghs=ghsR(p,i,j), radi=r(p,i,j,k))
                end do
                do p = pes(3), pee(3)
                  call sw_case(case_id, lon(p,i,j), lat(p,i,j), ghs=ghsB(p,i,j), radi=r(p,i,j,k))
                end do
                do p = pes(4), pee(4)
                  call sw_case(case_id, lon(p,i,j), lat(p,i,j), ghs=ghsT(p,i,j), radi=r(p,i,j,k))
                end do
                do p = pes(5), pee(5)
                  r(p,i,j,k) = earth_radius
                end do
                do p = pes(6), pee(6)
                  r(p,i,j,k) = earth_radius
                end do
                
                ! Set gravity, ghs and coriolis on cell quadrature points.
                do p = pqs, pqe
                  call sw_case(case_id, lon(p,i,j), lat(p,i,j), ghs=ghsQ(p,i,j), g=gravityQ(p,i,j,k), f=fQ(p,i,j), radi=r(p,i,j,k))
                end do
                
                ! Set radius on corner points.
                do p = pvs, pve
                  call sw_case(case_id, lon(p,i,j), lat(p,i,j), radi=r(p,i,j,k))
                enddo
                
                ! Set radius on ghost points.
                do p = 1, ghost(i,j)%ngp
                  call sw_case(case_id, ghost(i,j)%lon(p), ghost(i,j)%lat(p), radi=ghost(i,j)%r(p,k))
                enddo
              end do
            end do
          enddo
          !$OMP END PARALLEL DO
          
          call mesh%domain(idom)%set_hor_metrics
          
          ! Set fields
          !$OMP PARALLEL DO PRIVATE(p,u,v,ght,uc,vc,gd,Jh,JhQ,iAh) COLLAPSE(3)
          do k = kms, kme
            do j = jms, jme
              do i = ims, ime
                do p = pqs, pqe
                  iAh = mesh%get_dom_iAh(idom,p,i,j,k)
                  call sw_case(case_id, lon(p,i,j), lat(p,i,j), u, v, ght(p))
                  call spherev_to_contrav(iAh, u, v, uc(p), vc(p))
                end do
                
                ! Set initial values on cell quadrature points.
                do p = pqs, pqe
                  Jh = mesh%get_dom_Jh(idom,p,i,j,k)
                  qQ(p,i,j,k,1) = Jh * ghsQ(p,i,j)
                  qQ(p,i,j,k,2) = qQ(p,i,j,k,1) * uc(p)
                  qQ(p,i,j,k,3) = qQ(p,i,j,k,1) * vc(p)
                end do
                
                ! Sum quadrature points to get cell integrated average values.
                do p = pqs, pqe
                  JhQ(p-pqs+1) = mesh%get_dom_Jh(idom,p,i,j,k)
                enddo
                ghs(i,j) = mesh%domain(idom)%cell_quad(ghsQ(:,i,j))
                gd = ( ght - ghsQ(:,i,j) ) * JhQ
                uc = gd * uc
                vc = gd * vc
                q(i,j,1,1) = mesh%domain(idom)%cell_quad(gd)
                q(i,j,1,2) = mesh%domain(idom)%cell_quad(uc)
                q(i,j,1,3) = mesh%domain(idom)%cell_quad(vc)
              end do
            end do
          enddo
          !$OMP END PARALLEL DO
          end associate
        end if
      end do
      deallocate(gd, ght, uc, vc)
      end associate
      
    elseif(model_type==2)then
      if(case_id==0)then
        ! Real case
        call mesh%set_hor_metrics
        
        call topo_init(static)
        
        call vertical_init(static)
        
        call test_set_radius(static)
        
        call topo_deriv_init(static)
        
        call vertical_deriv_init(static)
        
        call test_unify_metric(static)

        call test_correct_metric(static)
        
        !call reconstruct_metric(static)
        
        case_omega = earth_omega
        call test_set_omega(static)
        
        call test_set_gravity(static)
        
        call test_set_damping(static)
        dvs = 4
        dve = 4
        
        high_field_file      = '../source_data/ERA-interim_high_2019_July.nc'
        surface_field_file   = '../source_data/ERA-interim_surface_2019_July.nc'
        invariant_field_file = '../source_data/era_interim_invariant.nc'
        itime = 45
        call met_get_data_era_interim(high_field_file, surface_field_file, invariant_field_file, itime)
        call met_interp_fields(state, static)
        
        ! Set Reference state
        do idom = 1,mesh%ndom
          if (mesh%domain(idom)%initialized) then
            state_ref%domain(idom)%q = 0
          endif
        enddo
        
      elseif(case_id>=1)then
        ! non-hydrostatic ideal case
        call mesh%set_hor_metrics
        
        associate (mesh => state%mesh)
          call mesh%get_params(npt=npt, pc=pc, pvs=pvs, pve=pve, pes=pes, pee=pee, pqs=pqs, pqe=pqe, pgs=pgs, pge=pge)
          do idom = 1, mesh%ndom
            if (mesh%domain(idom)%initialized) then
              associate (ghost    => mesh%domain(idom)%ghost      , &
                         lon      => mesh%domain(idom)%lon        , &
                         lat      => mesh%domain(idom)%lat        , &
                         r        => mesh%domain(idom)%r          , &
                         zs       => static%domain(idom)%zs       , &
                         zsC      => static%domain(idom)%zsC      , &
                         zsL      => static%domain(idom)%zsL      , &
                         zsR      => static%domain(idom)%zsR      , &
                         zsB      => static%domain(idom)%zsB      , &
                         zsT      => static%domain(idom)%zsT      , &
                         zsD      => static%domain(idom)%zsD      , &
                         zsU      => static%domain(idom)%zsU      , &
                         zsQ      => static%domain(idom)%zsQ      , &
                         gravity  => static%domain(idom)%gravity  , &
                         gravityQ => static%domain(idom)%gravityQ   )
                call mesh%get_params(idom=idom, ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme)
                ! Set static data
                !$OMP PARALLEL DO PRIVATE(p) COLLAPSE(3)
                do k = kms, kme
                  do j = jms, jme
                    do i = ims, ime
                      ! Set gravity, zs and coriolis on center point
                      p = pc
                      call nh_case(case_id, lon(pc,i,j), lat(pc,i,j), zs=zsC(i,j), radi=r(p,i,j,k), omg=case_omega, zt=max_height)
                      
                      ! Set zs on edge quadrature points.
                      do p = pes(1), pee(1)
                        call nh_case(case_id, lon(p,i,j), lat(p,i,j), zs=zsL(p,i,j), radi=r(p,i,j,k))
                      end do
                      do p = pes(2), pee(2)
                        call nh_case(case_id, lon(p,i,j), lat(p,i,j), zs=zsR(p,i,j), radi=r(p,i,j,k))
                      end do
                      do p = pes(3), pee(3)
                        call nh_case(case_id, lon(p,i,j), lat(p,i,j), zs=zsB(p,i,j), radi=r(p,i,j,k))
                      end do
                      do p = pes(4), pee(4)
                        call nh_case(case_id, lon(p,i,j), lat(p,i,j), zs=zsT(p,i,j), radi=r(p,i,j,k))
                      end do
                      do p = pes(5), pee(5)
                        call nh_case(case_id, lon(p,i,j), lat(p,i,j), zs=zsD(p,i,j), radi=r(p,i,j,k))
                      end do
                      do p = pes(6), pee(6)
                        call nh_case(case_id, lon(p,i,j), lat(p,i,j), zs=zsU(p,i,j), radi=r(p,i,j,k))
                      end do
                      
                      ! Set initial values on cell quadrature points.
                      do p = pqs, pqe
                        call nh_case(case_id, lon(p,i,j), lat(p,i,j), zs=zsQ(p,i,j), g=gravityQ(p,i,j,k), radi=r(p,i,j,k))
                      end do
                      gravity(i,j,k) = mesh%domain(idom)%cell_quad( gravityQ(pqs:pqe,i,j,k) )
                      zs(i,j) =  mesh%domain(idom)%cell_quad(zsQ(:,i,j))
                      
                      ! Set zs on corner points.
                      do p = pvs, pve
                        call nh_case(case_id, lon(p,i,j), lat(p,i,j), radi=r(p,i,j,k))
                      enddo
                      
                      ! Set zs on ghost points.
                      do p = 1, ghost(i,j)%ngp
                        call nh_case(case_id, ghost(i,j)%lon(p), ghost(i,j)%lat(p), zs=ghost(i,j)%zs(p), radi=ghost(i,j)%r(p,k))
                      enddo
                    end do
                  end do
                enddo
                !$OMP END PARALLEL DO
              end associate
            endif
          enddo
          
          call vertical_init(static)
          
          call topo_deriv_init(static)
          
          call vertical_deriv_init(static)
          
          call test_unify_metric(static)

          call test_correct_metric(static)
          
          call test_set_omega(static)
          dvs = 4
          dve = 4
          
          allocate(rho(pqs:pqe), pt(pqs:pqe), mr(pqs:pqe), uc(pqs:pqe), vc(pqs:pqe), wc(pqs:pqe))
          do idom = 1, mesh%ndom
            if (mesh%domain(idom)%initialized) then
              associate (lon      => mesh%domain(idom)%lon    , &
                         lat      => mesh%domain(idom)%lat    , &
                         Jab      => mesh%domain(idom)%Jab    , &
                         r        => mesh%domain(idom)%r      , &
                         z        => mesh%domain(idom)%z      , &
                         rdc      => static%domain(idom)%rdc  , &
                         q        => state%domain(idom)%q     , &
                         q_ref    => state_ref%domain(idom)%q   )
                call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
                !$OMP PARALLEL DO PRIVATE(p,iAh,rho,pt,mr,u,v,w,uc,vc,wc) COLLAPSE(3)
                do k = kds, kde
                  do j = jds, jde
                    do i = ids, ide
                      p = pc
                      call nh_case(case_id, lon(p,i,j), lat(p,i,j), z=z(p,i,j,k), rc=rdc(i,j,k))
                      
                      ! Set initial values on cell quadrature points.
                      do p = pqs, pqe
                        iAh = mesh%get_dom_iAh(idom,p,i,j,k)
                        call nh_case(case_id, lon(p,i,j), lat(p,i,j), z=z(p,i,j,k), rho=rho(p), pt=pt(p), mr=mr(p), u=u, v=v, w=w)
                        call spherev_to_contrav(iAh, u, v, uc(p), vc(p))
                        wc(p) = w
                      end do
                      q(i,j,k,1) = mesh%domain(idom)%cell_quad( Jab(pqs:pqe,i,j,k) * rho      ) / Jab(pc,i,j,k)
                      q(i,j,k,2) = mesh%domain(idom)%cell_quad( Jab(pqs:pqe,i,j,k) * rho * uc ) / Jab(pc,i,j,k)
                      q(i,j,k,3) = mesh%domain(idom)%cell_quad( Jab(pqs:pqe,i,j,k) * rho * vc ) / Jab(pc,i,j,k)
                      q(i,j,k,4) = mesh%domain(idom)%cell_quad( Jab(pqs:pqe,i,j,k) * rho * wc ) / Jab(pc,i,j,k)
                      q(i,j,k,5) = mesh%domain(idom)%cell_quad( Jab(pqs:pqe,i,j,k) * rho * pt ) / Jab(pc,i,j,k)
                      q(i,j,k,6) = mesh%domain(idom)%cell_quad( Jab(pqs:pqe,i,j,k) * rho * mr ) / Jab(pc,i,j,k)
                      
                      q_ref(i,j,k,:) = q(i,j,k,:)
                    end do
                  end do
                enddo
                !$OMP END PARALLEL DO
              end associate
            endif
          enddo
        end associate
      endif
    endif
    
    call fill_halo(state%array)
    call ghost_interp(state, diag)

  end subroutine test_setup
  
  subroutine test_set_radius(static)
    type(static_type), intent(inout) :: static

    integer idom, ims, ime, jms, jme, kms, kme
    integer       ids, ide, jds, jde, kds, kde
    integer npt, pc, pvs, pve, pes(6), pee(6), pqs, pqe, pgs, pge

    integer i, j, k, p
        
    call log_notice('Calculate radius on points.', pid=proc%id)
    associate (mesh => static%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%initialized) then
        call mesh%get_params(idom=idom, npt=npt, pqs=pqs, pqe=pqe,               &
                             ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme)
        associate( domain => mesh%domain(idom),      &
                   ghost  => mesh%domain(idom)%ghost )
        !$OMP PARALLEL DO PRIVATE(p) COLLAPSE(3)
        do k = kms, kme
          do j = jms, jme
            do i = ims, ime
              do p = 1, npt
                domain%r(p,i,j,k) = earth_radius + domain%z(p,i,j,k)
              enddo
              do p = 1, ghost(i,j)%ngp
                ghost(i,j)%r(p,k) = earth_radius + ghost(i,j)%z(p,k)
              enddo
            enddo
          enddo
        enddo
        !$OMP END PARALLEL DO
        end associate
      endif
    end do
    end associate
  end subroutine test_set_radius
  
  subroutine test_unify_metric(static)
    type(static_type), intent(inout) :: static
  
    integer idom, ims, ime, jms, jme, kms, kme
    integer npt, pc, pqs, pqe
    real(r_kind) :: Jh

    integer i, j, k, p
        
    call log_notice('Unify Metric.', pid=proc%id)
    ! Unify metric
    associate (mesh => static%mesh)
      do idom = 1, mesh%ndom
        if (mesh%domain(idom)%initialized) then
          call mesh%get_params(idom=idom, npt=npt, pc=pc, pqs=pqs, pqe=pqe,        &
                               ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme)
          associate( ghost => mesh%domain(idom)%ghost,&
                     lat   => mesh%domain(idom)%lat,  &
                     Jab   => mesh%domain(idom)%Jab,  &
                     Jv    => mesh%domain(idom)%Jv    )
            !$OMP PARALLEL DO PRIVATE(p,i,j,Jh) COLLAPSE(3)
            do k = kms, kme
              do j = jms, jme
                do i = ims, ime
                  do p = 1, npt
                    Jh = mesh%get_dom_Jh(idom,p,i,j,k)
                    Jab(p,i,j,k) = Jh * Jv(p,i,j,k)
                  enddo
      
                  Jab(pc,i,j,k) = mesh%domain(idom)%cell_quad(Jab(pqs:pqe,i,j,k))
                  Jv (pc,i,j,k) = mesh%domain(idom)%cell_quad(Jv (pqs:pqe,i,j,k))
                enddo
              enddo
            enddo
            !$OMP END PARALLEL DO
          end associate
        endif
      end do
    end associate
  
  end subroutine test_unify_metric

  subroutine test_correct_metric(static)
    type(static_type), intent(inout) :: static
  
    integer idom, ims, ime, jms, jme, kms, kme
    integer       ids, ide, jds, jde, kds, kde
    integer npt, pc, pvs, pve, pes(6), pee(6), pqs, pqe, pgs, pge
    real(r_kind) :: JabE(neq), JabQ(ncq)
    real(r_kind) :: iGvE(2,3:3,neq), iGvQ(2,3:3,ncq)
    real(r_kind) :: xi ,eta, z, r
    real(r_kind) :: dxideta, detadxi
    real(r_kind) :: Ra, Rb, Rr
    real(r_kind) :: area_old, area_new

    integer i, j, k, p, ip, iEdge, irow, icol
        
    call log_notice('Correct Metric.', pid=proc%id)
    associate (mesh => static%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%initialized) then
        call mesh%get_params(idom=idom, npt=npt, pc=pc, pqs=pqs, pqe=pqe, pes=pes, pee=pee, &
                             ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme)
        associate( ghost   => mesh%domain(idom)%ghost,    &
                   xeq     => mesh%domain(idom)%xeq,      &
                   lat     => mesh%domain(idom)%lat,      &
                   Jab     => mesh%domain(idom)%Jab,      &
                   Jh      => mesh%domain(idom)%Jh,       &
                   Jv      => mesh%domain(idom)%Jv,       &
                   iGV     => mesh%domain(idom)%iGV,      &
                   eta     => mesh%domain(idom)%eta,      &
                   deta    => mesh%domain(idom)%deta,     &
                   eta_min => mesh%domain(idom)%etamin,   & 
                   eta_max => mesh%domain(idom)%etamax,   & 
                   z_min   => mesh%domain(idom)%zmin,     &
                   z_max   => mesh%domain(idom)%zmax,     &
                   zs      => static%domain(idom)%zs,     &
                   zsC     => static%domain(idom)%zsC,    &
                   zsL     => static%domain(idom)%zsL,    &
                   zsR     => static%domain(idom)%zsR,    &
                   zsB     => static%domain(idom)%zsB,    &
                   zsT     => static%domain(idom)%zsT,    &
                   zsU     => static%domain(idom)%zsU,    &
                   zsD     => static%domain(idom)%zsD,    &
                   zsQ     => static%domain(idom)%zsQ,    &
                   dzsdxC  => static%domain(idom)%dzsdxC, &
                   dzsdxL  => static%domain(idom)%dzsdxL, &
                   dzsdxR  => static%domain(idom)%dzsdxR, &
                   dzsdxB  => static%domain(idom)%dzsdxB, &
                   dzsdxT  => static%domain(idom)%dzsdxT, &
                   dzsdxU  => static%domain(idom)%dzsdxU, &
                   dzsdxD  => static%domain(idom)%dzsdxD, &
                   dzsdxQ  => static%domain(idom)%dzsdxQ, &
                   dzsdy   => static%domain(idom)%dzsdy,  &
                   dzsdyC  => static%domain(idom)%dzsdyC, &
                   dzsdyL  => static%domain(idom)%dzsdyL, &
                   dzsdyR  => static%domain(idom)%dzsdyR, &
                   dzsdyB  => static%domain(idom)%dzsdyB, &
                   dzsdyT  => static%domain(idom)%dzsdyT, &
                   dzsdyU  => static%domain(idom)%dzsdyU, &
                   dzsdyD  => static%domain(idom)%dzsdyD, &
                   dzsdyQ  => static%domain(idom)%dzsdyQ  )
          !$OMP PARALLEL DO PRIVATE(p,i,j,iEdge,ip,r,JabE,JabQ,iGvE,iGvQ,area_new,area_old,xi,dxideta,detadxi,z,Ra,Rb,Rr) COLLAPSE(3)
          do k = kms, kme
            do j = jms, jme
              do i = ims, ime
                ! Left
                iEdge = 1
                do p = pes(iEdge), pee(iEdge)
                  do ip = 1, neq
                    call vertical_distribution(xi,dxideta,detadxi,eta(pes(5),k)+real(xeq(ip),r_kind)*deta,eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z,xi,zsL(p,i,j),z_min,z_max)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi,zsL(p,i,j),dzsdxL(p,i,j),dzsdyL(p,i,j),dxideta,detadxi,z_min,z_max)

                    r            = earth_radius + z
                    JabE(    ip) = Jh(p,i,j) * r * r * Rr
                    iGvE(1,3,ip) = -Ra / Rr
                    iGvE(2,3,ip) = -Rb / Rr
                  enddo
                  area_new = mesh%domain(idom)%edge_quad(JabE)
                  area_old = Jab(p,i,j,k)
                  Jab(p,i,j,k) = area_new / area_old * Jab(p,i,j,k)
                  do icol = 3,3
                    do irow = 1,2
                      iGv(irow,icol,p,i,j,k) = mesh%domain(idom)%edge_quad( iGvE(irow,icol,:) )
                    enddo
                  enddo

                enddo

                ! Right
                iEdge = 2
                do p = pes(iEdge), pee(iEdge)
                  do ip = 1, neq
                    call vertical_distribution(xi,dxideta,detadxi,eta(pes(5),k)+real(xeq(ip),r_kind)*deta,eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z,xi,zsR(p,i,j),z_min,z_max)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi,zsR(p,i,j),dzsdxR(p,i,j),dzsdyR(p,i,j),dxideta,detadxi,z_min,z_max)

                    r            = earth_radius + z
                    JabE(    ip) = Jh(p,i,j) * r * r * Rr
                    iGvE(1,3,ip) = -Ra / Rr
                    iGvE(2,3,ip) = -Rb / Rr
                  enddo
                  area_new = mesh%domain(idom)%edge_quad(JabE)
                  area_old = Jab(p,i,j,k)
                  Jab(p,i,j,k) = area_new / area_old * Jab(p,i,j,k)
                  do icol = 3,3
                    do irow = 1,2
                      iGv(irow,icol,p,i,j,k) = mesh%domain(idom)%edge_quad( iGvE(irow,icol,:) )
                    enddo
                  enddo

                enddo

                ! Bottom
                iEdge = 3
                do p = pes(iEdge), pee(iEdge)
                  do ip = 1, neq
                    call vertical_distribution(xi,dxideta,detadxi,eta(pes(5),k)+real(xeq(ip),r_kind)*deta,eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z,xi,zsB(p,i,j),z_min,z_max)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi,zsB(p,i,j),dzsdxB(p,i,j),dzsdyB(p,i,j),dxideta,detadxi,z_min,z_max)

                    r            = earth_radius + z
                    JabE(    ip) = Jh(p,i,j) * r * r * Rr
                    iGvE(1,3,ip) = -Ra / Rr
                    iGvE(2,3,ip) = -Rb / Rr
                  enddo
                  area_new = mesh%domain(idom)%edge_quad(JabE)
                  area_old = Jab(p,i,j,k)
                  Jab(p,i,j,k) = area_new / area_old * Jab(p,i,j,k)
                  do icol = 3,3
                    do irow = 1,2
                      iGv(irow,icol,p,i,j,k) = mesh%domain(idom)%edge_quad( iGvE(irow,icol,:) )
                    enddo
                  enddo

                enddo

                ! Top
                iEdge = 4
                do p = pes(iEdge), pee(iEdge)
                  do ip = 1, neq
                    call vertical_distribution(xi,dxideta,detadxi,eta(pes(5),k)+real(xeq(ip),r_kind)*deta,eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z,xi,zsT(p,i,j),z_min,z_max)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi,zsT(p,i,j),dzsdxT(p,i,j),dzsdyT(p,i,j),dxideta,detadxi,z_min,z_max)

                    r            = earth_radius + z
                    JabE(    ip) = Jh(p,i,j) * r * r * Rr
                    iGvE(1,3,ip) = -Ra / Rr
                    iGvE(2,3,ip) = -Rb / Rr
                  enddo
                  area_new = mesh%domain(idom)%edge_quad(JabE)
                  area_old = Jab(p,i,j,k)
                  Jab(p,i,j,k) = area_new / area_old * Jab(p,i,j,k)
                  do icol = 3,3
                    do irow = 1,2
                      iGv(irow,icol,p,i,j,k) = mesh%domain(idom)%edge_quad( iGvE(irow,icol,:) )
                    enddo
                  enddo

                enddo

                ! Down
                do p = pqs, pqe
                  ip = p - pqs + 1
                  call vertical_distribution(xi,dxideta,detadxi,eta(pes(5),k),eta_min,eta_max,z_min,z_max,m_coef)
                  call vertical_coordinate(z,xi,zsQ(p,i,j),z_min,z_max)
                  call vertical_coordinate_deriv(Ra,Rb,Rr,xi,zsQ(p,i,j),dzsdxQ(p,i,j),dzsdyQ(p,i,j),dxideta,detadxi,z_min,z_max)

                  r            = earth_radius + z
                  JabQ(    ip) = Jh(p,i,j) * r * r * Rr
                  iGvQ(1,3,ip) = -Ra / Rr
                  iGvQ(2,3,ip) = -Rb / Rr
                enddo
                area_new = mesh%domain(idom)%cell_quad(JabQ)
                area_old = mesh%domain(idom)%edge_quad_ver(Jab(pes(5):pee(5),i,j,k))
                Jab(pes(5):pee(5),i,j,k) = area_new / area_old * Jab(pes(5):pee(5),i,j,k)
                do p = pes(5), pee(5)
                  do icol = 3,3
                    do irow = 1,2
                      iGv(irow,icol,p,i,j,k) = mesh%domain(idom)%cell_quad( iGvQ(irow,icol,:) )
                    enddo
                  enddo
                enddo

                ! Up
                do p = pqs, pqe
                  ip = p - pqs + 1
                  call vertical_distribution(xi,dxideta,detadxi,eta(pes(6),k),eta_min,eta_max,z_min,z_max,m_coef)
                  call vertical_coordinate(z,xi,zsQ(p,i,j),z_min,z_max)
                  call vertical_coordinate_deriv(Ra,Rb,Rr,xi,zsQ(p,i,j),dzsdxQ(p,i,j),dzsdyQ(p,i,j),dxideta,detadxi,z_min,z_max)

                  r            = earth_radius + z
                  JabQ(    ip) = Jh(p,i,j) * r * r * Rr
                  iGvQ(1,3,ip) = -Ra / Rr
                  iGvQ(2,3,ip) = -Rb / Rr
                enddo
                area_new = mesh%domain(idom)%cell_quad(JabQ)
                area_old = mesh%domain(idom)%edge_quad_ver(Jab(pes(6):pee(6),i,j,k))
                Jab(pes(6):pee(6),i,j,k) = area_new / area_old * Jab(pes(6):pee(6),i,j,k)
                do p = pes(6), pee(6)
                  do icol = 3,3
                    do irow = 1,2
                      iGv(irow,icol,p,i,j,k) = mesh%domain(idom)%cell_quad( iGvQ(irow,icol,:) )
                    enddo
                  enddo
                enddo

                ! Quad
                do p = pqs, pqe
                  do ip = 1, neq
                    call vertical_distribution(xi,dxideta,detadxi,eta(pes(5),k)+real(xeq(ip),r_kind)*deta,eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z,xi,zsQ(p,i,j),z_min,z_max)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi,zsQ(p,i,j),dzsdxQ(p,i,j),dzsdyQ(p,i,j),dxideta,detadxi,z_min,z_max)

                    r            = earth_radius + z
                    JabE(    ip) = Jh(p,i,j) * r * r * Rr
                    iGvE(1,3,ip) = -Ra / Rr
                    iGvE(2,3,ip) = -Rb / Rr
                  enddo
                  area_new = mesh%domain(idom)%edge_quad(JabE)
                  area_old = Jab(p,i,j,k)
                  Jab(p,i,j,k) = area_new / area_old * Jab(p,i,j,k)
                  do icol = 3,3
                    do irow = 1,2
                      iGv(irow,icol,p,i,j,k) = mesh%domain(idom)%edge_quad( iGvE(irow,icol,:) )
                    enddo
                  enddo
                enddo

                Jab(pc,i,j,k) = mesh%domain(idom)%cell_quad(Jab(pqs:pqe,i,j,k))

                do icol = 3,3
                  do irow = 1,2
                    iGv(irow,icol,pc,i,j,k) = mesh%domain(idom)%cell_quad( iGv(irow,icol,pqs:pqe,i,j,k) )
                  enddo
                enddo

              enddo
            enddo
          enddo
          !$OMP END PARALLEL DO
        end associate
      endif
    end do
    end associate

  end subroutine test_correct_metric

  subroutine test_set_omega(static)
    type(static_type), intent(inout) :: static

    integer idom, ims, ime, jms, jme, kms, kme
    integer       ids, ide, jds, jde, kds, kde
    integer npt, pc, pvs, pve, pes(6), pee(6), pqs, pqe, pgs, pge

    real(8), dimension(3,3) :: iA
    real(8), dimension(3  ) :: omega
    integer i, j, k, p
        
    call log_notice('Set planet angular velocity.', pid=proc%id)
    ! Calculate earth omega on cubed sphere
    associate (mesh => static%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%initialized) then
        call mesh%get_params(idom=idom, npt=npt, pqs=pqs, pqe=pqe,               &
                             ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        associate( lat => mesh%domain(idom)%lat,&
                   omg => mesh%domain(idom)%omg )
          !$OMP PARALLEL DO PRIVATE(p,i,j,omega,iA) COLLAPSE(4)
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                do p = pqs, pqe
                  iA = 0
                  iA(1:2,1:2) = mesh%get_dom_iAh(idom,p,i,j,k)
                  iA(3,3) = 1
                  omega = case_omega * [0._r_kind, cos(lat(p,i,j)), sin(lat(p,i,j))]
                  omg(:,p,i,j,k) = matmul(iA, omega)
                enddo
              enddo
            enddo
          enddo
          !$OMP END PARALLEL DO
        end associate
      endif
    end do
    end associate

  end subroutine test_set_omega

  subroutine test_set_gravity(static)
    type(static_type), intent(inout) :: static
    
    integer idom, ims, ime, jms, jme, kms, kme
    integer pc, pqs, pqe
    integer i, j, k, p
    
    call log_notice('Set gravity acceleration.', pid=proc%id)
    ! Set gravity
    associate (mesh => static%mesh)
      call mesh%get_params(pc=pc, pqs=pqs, pqe=pqe)
      do idom = 1, mesh%ndom
        if(mesh%domain(idom)%initialized)then
          call mesh%get_params(idom=idom, ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme)
          associate( Jab => mesh%domain(idom)%Jab )
          !$OMP PARALLEL DO PRIVATE(p,i,j) COLLAPSE(3)
          do k = kms, kme
            do j = jms, jme
              do i = ims, ime
                do p = pqs, pqe
                  if(use_3d_gravity)then
                    static%domain(idom)%gravityQ(p,i,j,k) = earth_gravity * earth_radius**2 / mesh%domain(idom)%r(p,i,j,k)**2
                  else
                    static%domain(idom)%gravityQ(p,i,j,k) = earth_gravity
                  endif
                enddo
                static%domain(idom)%gravity(i,j,k) = mesh%domain(idom)%cell_quad( Jab(pqs:pqe,i,j,k) * static%domain(idom)%gravityQ(pqs:pqe,i,j,k) ) / Jab(pc,i,j,k)
              enddo
            enddo
          enddo
          !$OMP END PARALLEL DO
          end associate
        endif
      enddo
    end associate
  
  end subroutine test_set_gravity
  
  subroutine test_set_damping(static)
    type(static_type), intent(inout) :: static
    
    integer idom, ids, ide, jds, jde, kds, kde, pc, pqs, pqe, pes(6), pee(6)
    integer i, j, k, p
    real(r_kind) :: zd, zt
        
    call log_notice('Set w damping coefficients.', pid=proc%id)
    associate (mesh => static%mesh)
      call mesh%get_params(pqs=pqs, pqe=pqe)
      do idom = 1, mesh%ndom
        if(mesh%domain(idom)%active)then
          call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde, &
                               pc=pc, pqs=pqs, pqe=pqe, pes=pes, pee=pee)
          associate( z   => mesh%domain(idom)%z,    &
                     rdc => static%domain(idom)%rdc )
          !$OMP PARALLEL DO PRIVATE(p, zt, zd) COLLAPSE(3)
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                p = pc
                zt = z(p,i,j,kde)
                zd = zt - damp_thickness
                if( z(p,i,j,k) > zd )then
                  rdc(i,j,k) = damp_max_coef * ( ( z(p,i,j,k) - zd ) / ( zt - zd ) )**2 ! ( Li Xingliang, MWR, 2013 )
                else
                  rdc(i,j,k) = 0
                endif
              enddo
            enddo
          enddo
          !$OMP END PARALLEL DO
          end associate
        endif
      enddo
    end associate
  end subroutine test_set_damping
  
  subroutine test_adv_update_wind(case_id, time_in_seconds, diag)
    integer          , intent(in   ) :: case_id
    real(r_kind)     , intent(in   ) :: time_in_seconds
    type(diag_type  ), intent(inout) :: diag

    integer idom, ims, ime, jms, jme, kms, kme
    integer       ids, ide, jds, jde, kds, kde
    integer npt, pc, pvs, pve, pes(6), pee(6), pqs, pqe, pgs, pge
    real(8) u, v
    real(8), allocatable, dimension(:) :: gd, gh, uc, vc
    real(8), dimension(2,2) :: iAh
    real(8), dimension(3) :: omega
    integer itime
    integer i, j, k, p
    
    associate (mesh => diag%mesh)
      call mesh%get_params(npt=npt, pc=pc, pvs=pvs, pve=pve, pes=pes, pee=pee, pqs=pqs, pqe=pqe, pgs=pgs, pge=pge)
      allocate(uc(npt), vc(npt))
      do idom = 1, mesh%ndom
        if (mesh%domain(idom)%initialized) then
          associate (ghost => mesh%domain(idom)%ghost, &
                     lon   => mesh%domain(idom)%lon  , &
                     lat   => mesh%domain(idom)%lat  , &
                     uu    => diag%domain(idom)%uc   , &
                     ucL   => diag%domain(idom)%ucL  , &
                     ucR   => diag%domain(idom)%ucR  , &
                     ucB   => diag%domain(idom)%ucB  , &
                     ucT   => diag%domain(idom)%ucT  , &
                     ucQ   => diag%domain(idom)%ucQ  , &
                     vv    => diag%domain(idom)%vc   , &
                     vcL   => diag%domain(idom)%vcL  , &
                     vcR   => diag%domain(idom)%vcR  , &
                     vcB   => diag%domain(idom)%vcB  , &
                     vcT   => diag%domain(idom)%vcT  , &
                     vcQ   => diag%domain(idom)%vcQ  , &
                     qL    => diag%domain(idom)%qL   , &
                     qR    => diag%domain(idom)%qR   , &
                     qB    => diag%domain(idom)%qB   , &
                     qT    => diag%domain(idom)%qT   , &
                     qQ    => diag%domain(idom)%qQ     )
          call mesh%get_params(idom=idom, ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme)
          ! Set fields
          !$OMP PARALLEL DO PRIVATE(p,u,v,gh,uc,vc,gd,iAh) COLLAPSE(3)
          do k = kms, kme
            do j = jms, jme
              do i = ims, ime
                do p = 1, npt
                  iAh = mesh%get_dom_iAh(idom,p,i,j,k)
                  call adv_case(case_id, time_in_seconds, lon(p,i,j), lat(p,i,j), u, v)
                  call spherev_to_contrav(iAh, u, v, uc(p), vc(p))
                end do
                
                ! Set ghs on edge quadrature points.
                do p = pes(1), pee(1)
                  ucL(p,i,j,k) = uc(p)
                  vcL(p,i,j,k) = vc(p)
                end do
                do p = pes(2), pee(2)
                  ucR(p,i,j,k) = uc(p)
                  vcR(p,i,j,k) = vc(p)
                end do
                do p = pes(3), pee(3)
                  ucB(p,i,j,k) = uc(p)
                  vcB(p,i,j,k) = vc(p)
                end do
                do p = pes(4), pee(4)
                  ucT(p,i,j,k) = uc(p)
                  vcT(p,i,j,k) = vc(p)
                end do
                
                ! Set initial values on cell quadrature points.
                do p = pqs, pqe
                  ucQ(p,i,j,k) = uc(p)
                  vcQ(p,i,j,k) = vc(p)
                end do
                
                ! Sum quadrature points to get cell integrated average values.
                uu(i,j,k) = mesh%domain(idom)%cell_quad(ucQ(pqs:pqe,i,j,k))
                vv(i,j,k) = mesh%domain(idom)%cell_quad(vcQ(pqs:pqe,i,j,k))
              end do
            end do
          enddo
          !$OMP END PARALLEL DO
          end associate
        end if
      end do
      deallocate(uc, vc)
      end associate
  end subroutine test_adv_update_wind
  
end module test_mod
